import cv2 as cv
import sys
import os
from python_ai.common.xcommon import *
import matplotlib.pyplot as plt
import numpy as np
import time
import datetime


def my_show_img(img, title="no title", trans=None, **kwargs):
    global spn
    spn += 1
    plt.subplot(spr, spc, spn)
    if trans is not None:
        img = trans(img)
    plt.imshow(img, **kwargs)
    plt.axis('off')
    plt.title(title)


img_dir = '../../../../../large_data/pic/'
img_name = 'football.jpg'
img_path = os.path.join(img_dir, img_name)

spr = 2
spc = 3
spn = 0
plt.figure(figsize=[12, 6])

img = cv.imread(img_path, cv.IMREAD_GRAYSCALE)
f = np.fft.fft2(img)
print_numpy_ndarray_info(f, 'f')
# my_show_img(f, 'f', cmap='gray')  # TypeError: Image data of dtype complex128 cannot be converted to float
fshift = np.fft.fftshift(f)
print(f'fshift: dtype: {fshift.dtype}, shape: {fshift.shape}, min: {fshift.min()}, max: {fshift.max()}')
print_numpy_ndarray_info(fshift, 'fshift')
# my_show_img(fshift, 'fshit', cmap='gray')  # TypeError: Image data of dtype complex128 cannot be converted to float
# 这里构建振幅图的公式没学过
magnitude_spectrum = 20 * np.log(np.abs(fshift))
# plt.figure()
# plt.subplot(121), plt.imshow(img, cmap='gray')
# plt.title('Input Image'), plt.xticks([]), plt.yticks([])
my_show_img(img, 'ori', cmap='gray')

# plt.subplot(122), plt.imshow(magnitude_spectrum, cmap='gray')
# plt.title('Magnitude Spectrum'), plt.xticks([]), plt.yticks([])
my_show_img(magnitude_spectrum, 'spectrum', cmap='gray')

# plt.show()

# plt.figure()

fshift2 = fshift.copy()

# HPF
rows, cols = img.shape[0], img.shape[1]
crow, ccol = rows // 2, cols // 2
fshift[crow - 30:crow + 30, ccol - 30:ccol + 30] = 0
f_ishift = np.fft.ifftshift(fshift)
img_back = np.fft.ifft2(f_ishift)
# my_show_img(img_back, 'img_back', cmap='gray')  # TypeError: Image data of dtype complex128 cannot be converted to float

# 取绝对值
img_back = np.abs(img_back)
my_show_img(img_back, 'img_back', cmap='gray')
my_show_img(img_back, 'img_back in JET')

# LPF
mask = np.zeros((rows, cols), dtype=np.uint8)
mask[crow - 30:crow + 30, ccol - 30:ccol + 30] = 1
fshift2 *= mask
f_ishift2 = np.fft.ifftshift(fshift2)
img_back2 = np.fft.ifft2(f_ishift2)
img_back2 = np.abs(img_back2)
my_show_img(img_back2, 'img_back2', cmap='gray')
my_show_img(img_back2, 'img_back2 in JET')


# plt.subplot(131), plt.imshow(img, cmap='gray')
# plt.title('Input Image'), plt.xticks([]), plt.yticks([])
# plt.subplot(132), plt.imshow(img_back, cmap='gray')
# plt.title('Image after HPF'), plt.xticks([]), plt.yticks([])
# plt.subplot(133), plt.imshow(img_back)
# plt.title('Result in JET'), plt.xticks([]), plt.yticks([])
plt.show()
